# ---------- 服务端代码 server.py ----------
"""
安全文件传输服务器
功能：接收客户端加密文件，验证文件类型，解密存储
使用说明：需预先配置加密密钥，客户端需按协议格式发送数据
"""

import socket
import threading
import os
from cryptography.fernet import Fernet  # 导入Fernet对称加密模块


# region 密钥管理模块
def load_key():
    """
    加载或生成加密密钥
    安全策略：
    - 密钥文件不存在时自动生成新密钥（首次运行）
    - 密钥存储为二进制文件，需保证文件系统安全
    - 相同密钥才能保证加解密成功，需妥善保管密钥文件
    """
    key_file = "secret.key"
    # 密钥不存在时生成新密钥（首次运行）
    if not os.path.exists(key_file):
        key = Fernet.generate_key()  # 生成256位(32字节)安全密钥
        with open(key_file, "wb") as f:
            f.write(key)
        print(f"[*] 已生成新密钥文件 {key_file}")

    # 读取现有密钥（后续运行）
    with open(key_file, "rb") as f:
        return f.read()  # 返回bytes类型的密钥


# endregion


# 初始化加密模块 ------------------------------------------------------------
key = load_key()  # 加载密钥（首次运行会生成）
fernet = Fernet(key)  # 创建Fernet加密器实例


# region 客户端处理函数
def handle_client(client_socket):
    """
    处理客户端连接的完整生命周期
    协议规范：
    1. 客户端首先发送4字节文件名长度（大端序）
    2. 接着发送实际文件名
    3. 发送4字节加密数据长度（大端序）
    4. 最后发送加密数据
    安全机制：
    - 文件名规范化防止路径遍历攻击
    - 文件扩展名白名单验证
    - 数据长度头校验防止DoS攻击
    """
    try:
        # region 阶段1：接收文件名 --------------------------------------------------
        # 读取文件名长度头（固定4字节）
        name_header = client_socket.recv(4)
        if not name_header:  # 客户端提前断开
            print("[!] 客户端断开连接")
            return

        # 将4字节转换为整数（大端序，网络标准字节序）
        name_len = int.from_bytes(name_header, byteorder='big')

        # 分块接收文件名数据（防止大文件导致内存溢出）
        filename = b''
        while len(filename) < name_len:
            # 计算剩余需要接收的字节数
            remaining = name_len - len(filename)
            # 每次最多接收4096字节（平衡效率和内存使用）
            chunk = client_socket.recv(min(4096, remaining))
            if not chunk:  # 连接意外中断
                raise ConnectionError("文件传输中断")
            filename += chunk
        # endregion

        # region 文件名安全处理 ----------------------------------------------------
        # 解码为字符串并规范化路径（防止../等路径遍历攻击）
        save_name = os.path.basename(filename.decode('utf-8')).strip()

        # 空文件名检查
        if not save_name:
            raise ValueError("无效文件名")

        # 文件扩展名白名单验证
        ALLOWED_EXT = ['.txt', '.jpg', '.xlsx', '.pdf']  # 可根据需求扩展
        if not any(save_name.lower().endswith(ext) for ext in ALLOWED_EXT):
            raise ValueError(f"禁止的文件类型: {save_name}")

        print(f"[*] 客户端请求保存为: {save_name}") #按照用户要求的文件名保存文件
        # endregion

        # region 阶段2：接收加密数据 ------------------------------------------------
        # 读取数据长度头（同样使用4字节大端序）
        data_header = client_socket.recv(4)
        if not data_header:
            raise ConnectionError("数据头接收失败")

        data_len = int.from_bytes(data_header, byteorder='big')

        # 分块接收加密数据（避免一次性接收大文件导致内存耗尽）
        encrypted = b''
        while len(encrypted) < data_len:
            # 动态计算每次接收的块大小
            chunk_size = min(4096, data_len - len(encrypted))
            chunk = client_socket.recv(chunk_size)
            if not chunk:
                raise ConnectionError("数据接收中断")
            encrypted += chunk
        # endregion

        # region 阶段3：解密存储 ---------------------------------------------------
        # 文件存在性检查（实际生产环境应考虑版本管理）
        if os.path.exists(save_name):
            print(f"[!] 文件 {save_name} 已存在，将被覆盖")

        # 使用Fernet解密数据（自动验证完整性）
        decrypted = fernet.decrypt(encrypted)

        # 二进制写入文件（保持原始数据格式）
        with open(save_name, "wb") as f:
            f.write(decrypted)

        print(f"[✓] 成功保存 {len(decrypted)} 字节到 {save_name}")
        # endregion

    except Exception as e:
        # 统一异常处理（记录错误类型和具体信息）
        print(f"[!] 处理错误: {type(e).__name__}: {e}")
    finally:
        # 确保关闭客户端套接字（释放资源）
        client_socket.close()


# endregion


# region 服务器主程序
def start_server():
    """
    启动文件传输服务器
    网络配置：
    - 监听所有网络接口（0.0.0.0）
    - 使用TCP协议，端口8080
    并发处理：
    - 每个客户端连接创建独立线程处理
    - 支持同时处理多个客户端请求
    """
    host = '127.0.0.1'  # 生产环境建议使用具体IP
    port = 8080

    # 创建TCP套接字（AF_INET: IPv4, SOCK_STREAM: TCP）
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # 设置端口重用选项（快速重启服务器）
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        # 绑定地址和端口
        server.bind((host, port))
        # 设置最大挂起连接数（backlog）
        server.listen(5)
        print(f"[*] 服务端启动在 {host}:{port}")
        print(f"[*] 等待客户端连接...")

        # 主循环持续接受新连接
        while True:
            # 接受客户端连接（阻塞等待）
            client, addr = server.accept()
            print(f"[+] 接受来自 {addr[0]}:{addr[1]} 的连接")

            # 创建新线程处理客户端（实现并发）
            handler = threading.Thread(target=handle_client, args=(client,))
            handler.start()

    except KeyboardInterrupt:
        # 处理Ctrl+C信号（优雅关闭）
        print("\n[!] 服务器关闭中...")
    finally:
        # 确保关闭服务器套接字
        server.close()


# endregion


if __name__ == "__main__":
    # 启动服务器（程序入口）
    start_server()